import copy
import math
import torch
import torch.nn as nn
from torch.nn import functional as F


# 线性层
def Linear(inputdim, outputdim, bias=True, uniform=True):
    linear = nn.Linear(inputdim, outputdim, bias)
    if uniform:
        nn.init.xavier_uniform_(linear.weight)
    else:
        nn.init.xavier_normal_(linear.weight)
    if bias:
        nn.init.constant_(linear.bias, 0.0)
    return linear


# 模型堆叠多层
def clone(module, N):
    return nn.ModuleList([copy.deepcopy(module) for _ in range(N)])


# Transformer中 “Multi-head Attention” 部分
class MultiHeadAttention(nn.Module):
    def __init__(self, head_num, d_model, dropout=0.0, v=1, output=1, relu=0):
        super(MultiHeadAttention, self).__init__()
        assert d_model % head_num == 0
        self.d_k = d_model // head_num
        self.head = head_num
        self.output = output

        self.userelu = relu

        self.linears = clone(Linear(d_model, d_model), 2)
        if v:
            self.linears.append(Linear(d_model, d_model))
        if output:
            self.linears.append(Linear(d_model, d_model))

        self.v = v
        self.output = output

        self.attn = None
        self.dropout = None
        if dropout > 0:
            self.dropout = nn.Dropout(dropout)

    def attention(self, q, k, v, mask=None):
        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k)
        # when searching, target mask is not needed
        if mask is not None:
            # b 1 t -> b 1 1 t -> b head t t
            mask = mask.unsqueeze(1).expand_as(scores)
            scores.masked_fill_(mask == 0, -1e9)
        p_att = F.softmax(scores, -1)
        if self.dropout:
            p_att = self.dropout(p_att)
        return torch.matmul(p_att, v)

    def forward(self, query, key, value, mask=None):
        # q k v : B T H
        nbatches = query.size(0)
        # b head t dim
        if self.v:
            if self.userelu:
                query, key, value = [F.relu(l(x)).view(nbatches, -1, self.head, self.d_k).transpose(1, 2)
                                     for l, x in zip(self.linears, (query, key, value))]
            else:
                query, key, value = [l(x).view(nbatches, -1, self.head, self.d_k).transpose(1, 2)
                                     for l, x in zip(self.linears, (query, key, value))]
        else:
            if self.userelu:
                query, key = [F.relu(l(x)).view(nbatches, -1, self.head, self.d_k).transpose(1, 2)
                              for l, x in zip(self.linears, (query, key))]
            else:
                query, key = [l(x).view(nbatches, -1, self.head, self.d_k).transpose(1, 2)
                              for l, x in zip(self.linears, (query, key))]
            value = value.view(nbatches, -1, self.head, self.d_k).transpose(1, 2)

        x = self.attention(query, key, value, mask)
        x = x.transpose(1, 2).contiguous().view(nbatches, -1, self.head * self.d_k)
        if self.output:
            x = self.linears[-1](x)
        # returen b t dim
        return x


# Transformer中 “add&norm” 部分
class AddandNorm(nn.Module):
    def __init__(self, size, dropout):
        super(AddandNorm, self).__init__()
        self.norm = nn.LayerNorm(size)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, x_trans):
        return self.norm(x + self.dropout(x_trans))


# Transformer中 “FFN” 部分
class PositionwiseFeedForward(nn.Module):
    def __init__(self, d_model, d_ff, dropout=0.0):
        super(PositionwiseFeedForward, self).__init__()
        self.w_1 = Linear(d_model, d_ff)
        self.w_2 = Linear(d_ff, d_model)
        self.dropout = None
        if dropout > 0:
            self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        h = F.relu(self.w_1(x), inplace=True)
        if self.dropout:
            h = self.dropout(h)
        return self.w_2(h)


class Attention(nn.Module):
    def __init__(self, dim, dropout_number=-1):
        super().__init__()
        self.dim = dim
        if dropout_number != -1:
            self.dropout = nn.Dropout(dropout_number)
        else:
            self.dropout = None

    def forward(self, q, k, v, mask=None):

        score = torch.matmul(q, k.transpose(-1, -2)) / math.sqrt(self.dim)
        if mask is not None:
            mask = mask.unsqueeze(-1).expand_as(score)
            score.masked_fill_(mask == 0, -1e9)

        prob = nn.Softmax(dim=-1)(score)
        if self.dropout is not None:
            prob = self.dropout(prob)
        out = torch.matmul(prob, v)
        return out, prob